package com.levi.threadpool;

import com.levi.utils.MDCUtil;
import org.slf4j.MDC;

import java.util.concurrent.*;

/**
 * @author shentong
 * @since 2023/1/9 14:46
 * 自定义线程池，用于传递MDC中的数据
 */
public class ThreadPoolExecutorMdcWrapper extends ThreadPoolExecutor {
    public ThreadPoolExecutorMdcWrapper(int corePoolSize, int maximumPoolSize, long keepAliveTime, TimeUnit unit, BlockingQueue<Runnable> workQueue) {
        super(corePoolSize, maximumPoolSize, keepAliveTime, unit, workQueue);
    }

    public ThreadPoolExecutorMdcWrapper(int corePoolSize, int maximumPoolSize, long keepAliveTime, TimeUnit unit, BlockingQueue<Runnable> workQueue, ThreadFactory threadFactory) {
        super(corePoolSize, maximumPoolSize, keepAliveTime, unit, workQueue, threadFactory);
    }

    public ThreadPoolExecutorMdcWrapper(int corePoolSize, int maximumPoolSize, long keepAliveTime, TimeUnit unit, BlockingQueue<Runnable> workQueue, RejectedExecutionHandler handler) {
        super(corePoolSize, maximumPoolSize, keepAliveTime, unit, workQueue, handler);
    }

    public ThreadPoolExecutorMdcWrapper(int corePoolSize, int maximumPoolSize, long keepAliveTime, TimeUnit unit, BlockingQueue<Runnable> workQueue, ThreadFactory threadFactory, RejectedExecutionHandler handler) {
        super(corePoolSize, maximumPoolSize, keepAliveTime, unit, workQueue, threadFactory, handler);
    }

    @Override
    public void execute(Runnable runnable) {
        super.execute(MDCUtil.wrap(runnable,MDC.getCopyOfContextMap()));
    }

    @Override
    public Future<?> submit(Runnable runnable) {
        return super.submit(MDCUtil.wrap(runnable,MDC.getCopyOfContextMap()));
    }

    @Override
    public <T> Future<T> submit(Runnable runnable, T result) {
        return super.submit(MDCUtil.wrap(runnable,MDC.getCopyOfContextMap()), result);
    }

    @Override
    public <T> Future<T> submit(Callable<T> callable) {
        return super.submit(MDCUtil.wrap(callable,MDC.getCopyOfContextMap()));
    }
}
